!  Purpose:
!  Hierarchical Bayesian inversion using Markov Chain Monte Carlo 
!  
!  Inputs: see below inputs to be stored in .nc file
!   
!  Outputs:
!  Must create a folder called 'outputs' in the current directory where a file 'output_file.nc' will be created (./outputs/output_file.nc) 
!
!  Requires:
!  Netcdf file *input_file.nc*, stored in folder called 'inputs' in current directory (./inputs/input_file.nc)
!   
!  Netcdf & LAPACK libraries
!   
!  Modules Used: MCMC_utils.f90, MCMC_subroutines.f90
!  
!  Written by: A.Ganesan, University of Bristol, July 18 2014

!  Calculates z = Cy + e1 where diag(e1) = D (measurement uncertainty, nugget term) (assumed Gaussian)
!  y = Hx + e2 where corr(e2) = the Kronecker product of spatial (S) and temporal (T) correlation matrices (assumed Gaussian)
!  Samples through y to find values for missing observations
!  Solves for separate variances for time and space and modifies determinant to be det(SXT)=det(S)^nmeasuretotal*det(T)^numsites
!  Exploits properties of the kronecker product such that the matrix-vector multiplication C = Rinv*n0 can be rewritten as
!  	Rinv*n0 = (Ainv kron Binv)*n0 = Y = Binv*N0_arr*Ainv' where N0_arr = reshape(n0,numsites,nmeasuretotal) and C = reshape(Y,nmeasuretotal*numsites,1)
!  This code never only requires explicit computation of the Kronecker product, inverse and determinant once in the code

Program hierarchical_MCMC

  use MCMC_utils, only        : kron
  use MCMC_subroutines, only  : setup, compute_x, record_all, &
       compute_sigma_y, compute_sigma_ys, &
       compute_tau, compute_nu_rho, &
       compute_y, compute_ratios
  ! Use everything, for now...
  use MCMC_data
#ifdef __USE_CUDA
  use MCMC_c_bindings
#endif

  Implicit none

  integer                     :: It,xi,yi
  real                        :: t_start

  call system_clock(count_0,count_rate,count_max)
  t_start=real(count_0)/real(count_rate)

  ! *************** Open Netcdf input and output files *****************

  call open_files()

  ! ************** Read input file *************************

  call read_params()

  !!***************** Allocate space for each variable ***************************

  call allocate_mem()
#ifdef __USE_CUDA
  call allocate_gpu_mem(nmeasuremax,numsites,nmeasure,statesize,dim1,&
                            dim2,nmeasuretotal)
#endif

  ! **************** Read in variables from input netcdf file *********************

  call read_data()

  !****************** Hierarchical Bayesian inversion********************* 
  ! Set number of threads for multi-threaded routines
  !$ call OMP_SET_NUM_THREADS(numthreads)
#if defined(__USE_MKL)
  call MKL_SET_NUM_THREADS(numthreads)
#endif

  call setup()

#ifdef __USE_CUDA
  ! Initialise the GPU with data
  call setup_gpu(statesize,pdf_param1_current,pdf_param2_current,x,p0_x,&
                 p0_pdf_param1,p0_pdf_param2,stepsize_pdf_param1,&
                 stepsize_pdf_param2,pdf_param1,pdf_param2,&
                 pdf_param1_pdf,pdf_param2_pdf,stepsize,pdf_hyperparam1,&
                 pdf_hyperparam2,H,x_pdf)
#endif

  call system_clock(count_1,count_rate,count_max)
  t0=real(count_1)/real(count_rate)
  write(*,*) 'Total time for setup (s):', t0-t_start

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!MCMC MAIN LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  do It=1,nIt+burn_in

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!X LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!	  

    ! Parameters here are the true inouts for compute_x
#ifdef __USE_CUDA
     call compute_x_gpu(nmeasuremax,nmeasuretotal,numsites,n0T,n0,&
                       Tinv_current,Sinv_current, it .gt. burn_in)
#else
     do xi=1,statesize
        call compute_x(stepsize_pdf_param1(xi),pdf_param1(xi),pdf_param1_pdf(xi),&
    pdf_param1_current(xi),stepsize_pdf_param2(xi),pdf_param2(xi),pdf_param2_pdf(xi),&
    pdf_param2_current(xi),stepsize(xi),H(:,xi),pdf_hyperparam1(xi),dy,&
    pdf_hyperparam2(xi),x(xi),x_pdf(xi),n1,n0,nmeasuretotal,nmeasuremax,&
    numsites,aa,Tinv_current,dum3,Sinv_current,bb,C,p0_x(xi),&
    p0_pdf_param1(xi),p0_pdf_param2(xi),n0T,reject,reject_vector(xi),it .gt. burn_in)
     enddo
#endif
#ifdef __USE_CUDA
     !Copy back the data that we need to record
     call copy_data_back(statesize,pdf_param1_current,pdf_param2_current,x,reject)
#endif

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!SIGMA_Y LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     do yi=1,dim2

        call compute_sigma_y(sigma_y_current(yi),sigma_y(yi),sigma_y_hyperparam(yi),&
    sigma_y_pdf(yi),stepsize_sigma_y(yi),Tinv_new,Tinv_current,T_indices(:,yi),&
    nmeasuretotal,numsites,nmeasuremax,aa,bb,dum3,Sinv_current,num_T_indices(yi),&
    detval_T_current,detval_S_current,p0_sigma_y(yi),n0T,n0,C,&
    detval_current,reject_sigma_y,reject_sigma_y_vector(yi),it .gt. burn_in)

     enddo

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!SIGMA_YS LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     do yi=1,numsites

        call compute_sigma_ys(yi,sigma_ys_current(yi),sigma_ys(yi),&
    sigma_ys_hyperparam(yi),sigma_ys_pdf(yi),stepsize_sigma_ys(yi),Sinv_new,&
    Sinv_current,nmeasuretotal,numsites,nmeasuremax,aa,bb,dum3,Tinv_current,&
    detval_S_current,detval_T_current,p0_sigma_ys(yi),n0T,n0,C,&
    detval_current,reject_sigma_ys,reject_sigma_ys_vector(yi),it .gt. burn_in)

     enddo

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!TAU LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     call compute_tau(stepsize_tau,tau_current,reject_tau,tau_hyperparam1,&
    tau_hyperparam2,tau_pdf,dim2,y_error_t,sigma_y_current,T_indices,&
    nmeasuretotal,autocorr_vec,aa,bb,numsites,nmeasuremax,&
    detval_S_current,p0_tau,detval_current,n0T,Tinv_current,&
    detval_T_current,n0,Sinv_current,deltatime,T_new,Tinv_new,dum3,&
    C,it .gt. burn_in)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!NU and RHO LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     call compute_nu_rho(stepsize_nu,stepsize_rho,nu_current,rho_current,&
    reject_nu,nu_hyperparam1,nu_hyperparam2,nu_pdf,rho_hyperparam1,rho_hyperparam2,rho_pdf,&
    arg,distance_double,sigma_ys_current,S_new,distance,numsites,Sinv_new,nmeasuretotal,&
    nmeasuremax,aa,bb,dum3,C,n0,n0T,p0_nu,p0_rho,detval_T_current,detval_current,&
    detval_S_current,Tinv_current,Sinv_current,it .gt. burn_in)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!Y LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

     ! create full covariance matrix R based on current state of S and T

     if (kron_flag==0) then
        call kron(Rinv_current, Sinv_current, Tinv_current)
     elseif (kron_flag==1) then
        call kron(Rinv_current, Tinv_current, Sinv_current)
     endif

     ! update n1 with most recent data
     n1 = n0
     do yi=1,nmeasuremax
        call compute_y(yi,y_current,stepsize_y(yi),m1,z,nmeasuremax,&
                       dum2,Dinv,n1,n0,dumx,dumy,Rinv_current,n0T,m0T,m1T,reject_y,&
                       reject_y_vector(yi),timeindex(yi),it .gt. burn_in)
     enddo

     if(it .gt. burn_in) then
        call record_all(It)
     endif

  end do

  call system_clock(count_1,count_rate,count_max)
  write(*,*) 'Total time for main loop (s):', real(count_1)/real(count_rate)-t0
  t0 = real(count_1)/real(count_rate)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!END MCMC MAIN LOOP!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  ! Compute acceptance ratios

#ifdef __USE_CUDA
  call copy_final_data_back(statesize,reject_vector)
#endif

  call compute_ratios()

  write(*,*) 'Finished HB_MCMC'

#ifdef __USE_CUDA
  call free_gpu_mem()
#endif

  ! ********************************** Write output to netcdf **************************************

  call write_data()

  !!************************************************************************************************************

  ! All arrays that are not explicitly saved will be deallocated upon exit

  call system_clock(count_1,count_rate,count_max)
  write(*,*) 'Total time for cleanup (s):', real(count_1)/real(count_rate)-t0
  write(*,*) 'Total time elapsed (s):', real(count_1)/real(count_rate)-t_start

  !!*****************************************************************************
  !!*****************************************************************************  
end program hierarchical_MCMC

